{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# SafeDrug Model Training on MIMIC-III Dataset\n",
        "\n",
        "Train the SafeDrug model for medication recommendation on the MIMIC-III dataset.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/Caskroom/miniforge/base/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "No config path provided, using default config\n",
            "Initializing mimic3 dataset from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III (dev mode: True)\n",
            "Scanning table: patients from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PATIENTS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: admissions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: icustays from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ICUSTAYS.csv\n",
            "Some column names were converted to lowercase\n",
            "Scanning table: diagnoses_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/DIAGNOSES_ICD.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Scanning table: procedures_icd from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PROCEDURES_ICD.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Scanning table: prescriptions from https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/PRESCRIPTIONS.csv\n",
            "Some column names were converted to lowercase\n",
            "Joining with table: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv.gz\n",
            "Original path does not exist. Using alternative: https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III/ADMISSIONS.csv\n",
            "Collecting global event dataframe...\n",
            "Dev mode enabled: limiting to 1000 patients\n",
            "Collected dataframe with shape: (42805, 49)\n",
            "Dataset: mimic3\n",
            "Dev mode: True\n",
            "Number of patients: 1000\n",
            "Number of events: 42805\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.datasets import MIMIC3Dataset\n",
        "\n",
        "dataset = MIMIC3Dataset(\n",
        "    root=\"https://storage.googleapis.com/pyhealth/Synthetic_MIMIC-III\",\n",
        "    tables=[\"DIAGNOSES_ICD\", \"PROCEDURES_ICD\", \"PRESCRIPTIONS\"],\n",
        "    dev=True,\n",
        ")\n",
        "dataset.stats()\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Set Drug Recommendation Task\n",
        "\n",
        "Use the `DrugRecommendationMIMIC3` task function which creates samples with conditions, procedures, and atc-3 codes (drugs).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Setting task DrugRecommendationMIMIC3 for mimic3 base dataset...\n",
            "Generating samples with 4 worker(s)...\n",
            "Generating samples for DrugRecommendationMIMIC3 with 4 workers\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Collecting samples for DrugRecommendationMIMIC3 from 4 workers: 100%|██████████| 1000/1000 [00:00<00:00, 1420.70it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Label drugs vocab: {'*NF*': 0, '1/2 ': 1, 'Acar': 2, 'Acet': 3, 'Acyc': 4, 'Albu': 5, 'Allo': 6, 'Alpr': 7, 'Alte': 8, 'Alum': 9, 'Amin': 10, 'Amio': 11, 'Amlo': 12, 'Amph': 13, 'Ampi': 14, 'Arga': 15, 'Arti': 16, 'Aspi': 17, 'Aten': 18, 'Ator': 19, 'Atro': 20, 'Augm': 21, 'Azit': 22, 'Baci': 23, 'Bisa': 24, 'Bume': 25, 'BusP': 26, 'Busp': 27, 'Calc': 28, 'Capt': 29, 'Carv': 30, 'Cefa': 31, 'Cefe': 32, 'Ceft': 33, 'Ceph': 34, 'Chlo': 35, 'Cipr': 36, 'Cisa': 37, 'Cita': 38, 'Clop': 39, 'Clot': 40, 'Cosy': 41, 'D10W': 42, 'D5 1': 43, 'D5NS': 44, 'D5W': 45, 'D5W ': 46, 'Daki': 47, 'Dapt': 48, 'Desm': 49, 'Dexa': 50, 'Dext': 51, 'Diaz': 52, 'Digo': 53, 'Dilt': 54, 'Diov': 55, 'Diph': 56, 'Diso': 57, 'Docu': 58, 'Dola': 59, 'DopA': 60, 'Doxa': 61, 'Enox': 62, 'Epin': 63, 'Eryt': 64, 'Famo': 65, 'Fat ': 66, 'Fent': 67, 'Ferr': 68, 'Fexo': 69, 'Filg': 70, 'Fina': 71, 'Fluc': 72, 'Flud': 73, 'Fluo': 74, 'Flut': 75, 'FoLI': 76, 'Foli': 77, 'Fosp': 78, 'Furo': 79, 'Gaba': 80, 'Gemf': 81, 'Gent': 82, 'Glip': 83, 'Gluc': 84, 'Glyc': 85, 'Guai': 86, 'HEPA': 87, 'HYDR': 88, 'Halo': 89, 'Hepa': 90, 'Humu': 91, 'Hydr': 92, 'Ibup': 93, 'Infl': 94, 'Insu': 95, 'Ipra': 96, 'Isos': 97, 'Kaop': 98, 'Keto': 99, 'LR': 100, 'Labe': 101, 'Lact': 102, 'Lans': 103, 'Lata': 104, 'LeVE': 105, 'Leva': 106, 'Leve': 107, 'Levo': 108, 'Lido': 109, 'Line': 110, 'Lisi': 111, 'Lora': 112, 'Losa': 113, 'Magn': 114, 'Mepe': 115, 'Mero': 116, 'Mesa': 117, 'MetR': 118, 'Meth': 119, 'Meto': 120, 'Metr': 121, 'Mico': 122, 'Mida': 123, 'Milk': 124, 'Milr': 125, 'Morp': 126, 'Mult': 127, 'Myco': 128, 'NEO*': 129, 'NIFE': 130, 'NORe': 131, 'NS': 132, 'NS (': 133, 'Neos': 134, 'Nesi': 135, 'Neut': 136, 'NiCA': 137, 'Nico': 138, 'Nife': 139, 'Nitr': 140, 'Nore': 141, 'Nyst': 142, 'Octr': 143, 'Olan': 144, 'Onda': 145, 'Oxyc': 146, 'Pant': 147, 'Paro': 148, 'Pent': 149, 'Phen': 150, 'Phyt': 151, 'Pipe': 152, 'Pneu': 153, 'Pota': 154, 'Prav': 155, 'Pred': 156, 'Preg': 157, 'Proc': 158, 'Prom': 159, 'Prop': 160, 'Prot': 161, 'Pyri': 162, 'Quet': 163, 'Quin': 164, 'Qvar': 165, 'Race': 166, 'Rani': 167, 'Risp': 168, 'Scop': 169, 'Senn': 170, 'Sert': 171, 'Seve': 172, 'Sime': 173, 'Simv': 174, 'Siro': 175, 'Sodi': 176, 'Spir': 177, 'Succ': 178, 'Sucr': 179, 'Sulf': 180, 'Tacr': 181, 'Tams': 182, 'Tema': 183, 'Thia': 184, 'Topi': 185, 'TraM': 186, 'Trav': 187, 'Tria': 188, 'Tube': 189, 'Tuck': 190, 'Urso': 191, 'Vanc': 192, 'Vaso': 193, 'Vera': 194, 'Vita': 195, 'Warf': 196, 'Zofr': 197, 'Zolp': 198, 'traM': 199, 'traZ': 200}\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Processing samples: 100%|██████████| 56/56 [00:00<00:00, 932.10it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Generated 56 samples for task DrugRecommendationMIMIC3\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sample Dataset Statistics:\n",
            "\t- Dataset: mimic3\n",
            "\t- Task: <pyhealth.tasks.drug_recommendation.DrugRecommendationMIMIC3 object at 0x133427b90>\n",
            "\t- Number of samples: 56\n",
            "\n",
            "First sample structure:\n",
            "Patient ID: 27333\n",
            "Number of visits: 1\n",
            "Sample conditions (first visit): tensor([1, 2, 3, 4, 5])...\n",
            "Sample procedures (first visit): tensor([1, 2, 3, 4, 5])...\n",
            "Sample drugs (target): tensor([0., 0., 0., 1., 0., 1., 0., 1., 1., 0.])...\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.tasks import DrugRecommendationMIMIC3\n",
        "\n",
        "task = DrugRecommendationMIMIC3()\n",
        "samples = dataset.set_task(task, num_workers=4)\n",
        "\n",
        "print(f\"Sample Dataset Statistics:\")\n",
        "print(f\"\\t- Dataset: {samples.dataset_name}\")\n",
        "print(f\"\\t- Task: {samples.task_name}\")\n",
        "print(f\"\\t- Number of samples: {len(samples)}\")\n",
        "\n",
        "print(\"\\nFirst sample structure:\")\n",
        "print(f\"Patient ID: {samples.samples[0]['patient_id']}\")\n",
        "print(f\"Number of visits: {len(samples.samples[0]['conditions'])}\")\n",
        "print(f\"Sample conditions (first visit): {samples.samples[0]['conditions'][0][:5]}...\")\n",
        "print(f\"Sample procedures (first visit): {samples.samples[0]['procedures'][0][:5]}...\")\n",
        "print(f\"Sample drugs (target): {samples.samples[0]['drugs'][:10]}...\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Split Dataset and Create Data Loaders\n",
        "\n",
        "Split the dataset by patient to ensure no data leakage between train/validation/test sets.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train samples: 38\n",
            "Validation samples: 4\n",
            "Test samples: 14\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.datasets import split_by_patient, get_dataloader\n",
        "\n",
        "train_dataset, val_dataset, test_dataset = split_by_patient(\n",
        "    samples, ratios=[0.7, 0.1, 0.2]\n",
        ")\n",
        "\n",
        "print(f\"Train samples: {len(train_dataset)}\")\n",
        "print(f\"Validation samples: {len(val_dataset)}\")\n",
        "print(f\"Test samples: {len(test_dataset)}\")\n",
        "\n",
        "train_dataloader = get_dataloader(train_dataset, batch_size=32, shuffle=True)\n",
        "val_dataloader = get_dataloader(val_dataset, batch_size=32, shuffle=False)\n",
        "test_dataloader = get_dataloader(test_dataset, batch_size=32, shuffle=False)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialize SafeDrug Model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/arjunchatterjee/PyHealth/pyhealth/sampler/sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
            "  import pkg_resources\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "SafeDrug(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(314, 128)\n",
            "    (procedures): Embedding(88, 128)\n",
            "    (drugs_hist): Embedding(166, 128)\n",
            "  ))\n",
            "  (cond_rnn): GRU(128, 128, batch_first=True)\n",
            "  (proc_rnn): GRU(128, 128, batch_first=True)\n",
            "  (query): Sequential(\n",
            "    (0): ReLU()\n",
            "    (1): Linear(in_features=256, out_features=128, bias=True)\n",
            "  )\n",
            "  (safedrug): SafeDrugLayer(\n",
            "    (bipartite_transform): Linear(in_features=128, out_features=0, bias=True)\n",
            "    (bipartite_output): Linear(in_features=0, out_features=201, bias=True)\n",
            "    (mpnn): MolecularGraphNeuralNetwork(\n",
            "      (embed_fingerprint): Embedding(1, 128)\n",
            "      (W_fingerprint): ModuleList(\n",
            "        (0-1): 2 x Linear(in_features=128, out_features=128, bias=True)\n",
            "      )\n",
            "    )\n",
            "    (mpnn_output): Linear(in_features=201, out_features=201, bias=True)\n",
            "    (mpnn_layernorm): LayerNorm((201,), eps=1e-05, elementwise_affine=True)\n",
            "    (test): Linear(in_features=128, out_features=201, bias=True)\n",
            "    (loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/Caskroom/miniforge/base/lib/python3.12/site-packages/torch/nn/init.py:452: UserWarning: Initializing zero-element tensors is a no-op\n",
            "  warnings.warn(\"Initializing zero-element tensors is a no-op\")\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.models import SafeDrug\n",
        "\n",
        "model = SafeDrug(\n",
        "    dataset=samples,\n",
        "    embedding_dim=128,\n",
        "    hidden_dim=128,\n",
        "    num_layers=1,\n",
        "    dropout=0.5,\n",
        ")\n",
        "\n",
        "print(model)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialize Trainer\n",
        "\n",
        "We use jaccard similarity, f1 score, pr_auc, ddi score.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "SafeDrug(\n",
            "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
            "    (conditions): Embedding(314, 128)\n",
            "    (procedures): Embedding(88, 128)\n",
            "    (drugs_hist): Embedding(166, 128)\n",
            "  ))\n",
            "  (cond_rnn): GRU(128, 128, batch_first=True)\n",
            "  (proc_rnn): GRU(128, 128, batch_first=True)\n",
            "  (query): Sequential(\n",
            "    (0): ReLU()\n",
            "    (1): Linear(in_features=256, out_features=128, bias=True)\n",
            "  )\n",
            "  (safedrug): SafeDrugLayer(\n",
            "    (bipartite_transform): Linear(in_features=128, out_features=0, bias=True)\n",
            "    (bipartite_output): Linear(in_features=0, out_features=201, bias=True)\n",
            "    (mpnn): MolecularGraphNeuralNetwork(\n",
            "      (embed_fingerprint): Embedding(1, 128)\n",
            "      (W_fingerprint): ModuleList(\n",
            "        (0-1): 2 x Linear(in_features=128, out_features=128, bias=True)\n",
            "      )\n",
            "    )\n",
            "    (mpnn_output): Linear(in_features=201, out_features=201, bias=True)\n",
            "    (mpnn_layernorm): LayerNorm((201,), eps=1e-05, elementwise_affine=True)\n",
            "    (test): Linear(in_features=128, out_features=201, bias=True)\n",
            "    (loss_fn): BCEWithLogitsLoss()\n",
            "  )\n",
            ")\n",
            "Metrics: ['jaccard_samples', 'f1_samples', 'pr_auc_samples', 'ddi']\n",
            "Device: cpu\n",
            "\n",
            "Baseline performance before training:\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Evaluation: 100%|██████████| 1/1 [00:03<00:00,  3.38s/it]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'jaccard_samples': 0.06751954513148543, 'f1_samples': 0.11837435592508397, 'pr_auc_samples': 0.06751954513148543, 'ddi_score': 0.0, 'loss': 0.6931472420692444}\n"
          ]
        }
      ],
      "source": [
        "from pyhealth.trainer import Trainer\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    metrics=[\"jaccard_samples\", \"f1_samples\", \"pr_auc_samples\", \"ddi\"],\n",
        ")\n",
        "\n",
        "print(\"Baseline performance before training:\")\n",
        "baseline_results = trainer.evaluate(test_dataloader)\n",
        "print(baseline_results)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train the Model\n",
        "\n",
        "Train the model for a few epochs. I used 5 epochs here. Might need to train for more epochs in prod.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Training:\n",
            "Batch size: 32\n",
            "Optimizer: <class 'torch.optim.adam.Adam'>\n",
            "Optimizer params: {'lr': 0.0001}\n",
            "Weight decay: 0.0\n",
            "Max grad norm: None\n",
            "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x1334c8b00>\n",
            "Monitor: pr_auc_samples\n",
            "Monitor criterion: max\n",
            "Epochs: 5\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 0 / 5: 100%|██████████| 2/2 [00:04<00:00,  2.36s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-0, step-2 ---\n",
            "loss: 0.6931\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 212.49it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-0, step-2 ---\n",
            "jaccard_samples: 0.1530\n",
            "f1_samples: 0.2510\n",
            "pr_auc_samples: 0.2577\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6931\n",
            "New best pr_auc_samples score (0.2577) at epoch-0, step-2\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 1 / 5: 100%|██████████| 2/2 [00:00<00:00, 40.41it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-1, step-4 ---\n",
            "loss: 0.6931\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 202.72it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-1, step-4 ---\n",
            "jaccard_samples: 0.1530\n",
            "f1_samples: 0.2510\n",
            "pr_auc_samples: 0.2659\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6930\n",
            "New best pr_auc_samples score (0.2659) at epoch-1, step-4\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 2 / 5: 100%|██████████| 2/2 [00:00<00:00, 45.36it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-2, step-6 ---\n",
            "loss: 0.6930\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 123.52it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-2, step-6 ---\n",
            "jaccard_samples: 0.1530\n",
            "f1_samples: 0.2510\n",
            "pr_auc_samples: 0.2776\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6930\n",
            "New best pr_auc_samples score (0.2776) at epoch-2, step-6\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 3 / 5: 100%|██████████| 2/2 [00:00<00:00, 44.84it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-3, step-8 ---\n",
            "loss: 0.6929\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 157.26it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-3, step-8 ---\n",
            "jaccard_samples: 0.1530\n",
            "f1_samples: 0.2510\n",
            "pr_auc_samples: 0.2989\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6929\n",
            "New best pr_auc_samples score (0.2989) at epoch-3, step-8\n",
            "\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Epoch 4 / 5: 100%|██████████| 2/2 [00:00<00:00, 41.89it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Train epoch-4, step-10 ---\n",
            "loss: 0.6928\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n",
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 185.05it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--- Eval epoch-4, step-10 ---\n",
            "jaccard_samples: 0.1530\n",
            "f1_samples: 0.2510\n",
            "pr_auc_samples: 0.3378\n",
            "ddi_score: 0.0000\n",
            "loss: 0.6928\n",
            "New best pr_auc_samples score (0.3378) at epoch-4, step-10\n",
            "Loaded best model\n"
          ]
        }
      ],
      "source": [
        "trainer.train(\n",
        "    train_dataloader=train_dataloader,\n",
        "    val_dataloader=val_dataloader,\n",
        "    epochs=5,\n",
        "    monitor=\"pr_auc_samples\",\n",
        "    optimizer_params={\"lr\": 1e-4},\n",
        ")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Evaluate on Test Set\n",
        "\n",
        "Evaluate the trained model on the test set to see final performance metrics.\n",
        "\n",
        "I get DDI = 0. That is abnormally low. (In a good way, I guess?)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 63.79it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Final test set performance:\n",
            "{'jaccard_samples': 0.06751954513148543, 'f1_samples': 0.11837435592508397, 'pr_auc_samples': 0.19516980311232707, 'ddi_score': 0.0, 'loss': 0.6927605271339417}\n",
            "\n",
            "Key Metrics:\n",
            "  PR-AUC: 0.1952\n",
            "  F1 Score: 0.1184\n",
            "  Jaccard: 0.0675\n",
            "  DDI Rate: 0.0000 (lower is better)\n"
          ]
        }
      ],
      "source": [
        "test_results = trainer.evaluate(test_dataloader)\n",
        "print(\"Final test set performance:\")\n",
        "print(test_results)\n",
        "\n",
        "print(f\"\\nKey Metrics:\")\n",
        "print(f\"  PR-AUC: {test_results.get('pr_auc_samples', 'N/A'):.4f}\")\n",
        "print(f\"  F1 Score: {test_results.get('f1_samples', 'N/A'):.4f}\")\n",
        "print(f\"  Jaccard: {test_results.get('jaccard_samples', 'N/A'):.4f}\")\n",
        "print(f\"  DDI Rate: {test_results.get('ddi_score', 'N/A'):.4f} (lower is better)\")\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "base",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
